import asyncio
import datetime
import json
import logging

import websockets

from model.bitget_depth import BitgetDepth
from model.bitget_ticker import BitgetTicker
import service.hedging_service as hedging_service


class WebSocketClient:
    def __init__(self, uri):
        self.uri = uri
        self.websocket = None
        self.keep_alive_task = None  # 保存任务的引用

    async def connect(self):
        while True:
            try:
                async with websockets.connect(self.uri, ping_interval=10, ping_timeout=5) as websocket:
                    logging.info("Connected")
                    # mix_book5 = {
                    #     "op": "subscribe",
                    #     "args": [
                    #         {
                    #             "instType": "USDT-FUTURES",
                    #             "channel": "books15",
                    #             "instId": "BTCUSDT"
                    #         }, {
                    #             "instType": "USDT-FUTURES",
                    #             "channel": "books15",
                    #             "instId": "ETHUSDT"
                    #         }
                    #     ]
                    # }
                    futures_ticker = {
                        "op": "subscribe",
                        "args": [
                            {
                                "instType": "USDT-FUTURES",
                                "channel": "ticker",
                                "instId": "BTCUSDT"
                            }, {
                                "instType": "USDT-FUTURES",
                                "channel": "ticker",
                                "instId": "ETHUSDT"
                            }
                        ]
                    }
                    futures_candles = hedging_service.get_subscribe_candles()
                    logging.info(futures_candles)
                    # futures_candles = {
                    #     "op": "subscribe",
                    #     "args": [
                    #         {
                    #             "instType": "USDT-FUTURES",
                    #             "channel": "candle1m",
                    #             "instId": "BTCUSDT"
                    #         }, {
                    #             "instType": "USDT-FUTURES",
                    #             "channel": "candle1m",
                    #             "instId": "ETHUSDT"
                    #         }
                    #     ]
                    # }
                    self.websocket = websocket
                    # await websocket.send(json.dumps(futures_ticker))
                    await websocket.send(json.dumps(futures_candles))
                    # 启动 keep_alive 任务，并存储任务的引用
                    if self.keep_alive_task is None or self.keep_alive_task.done():
                        self.keep_alive_task = asyncio.create_task(self.keep_alive())
                    await self.receive_message()
            except Exception as e:
                logging.info(f"Connection failed: {e}")
                await asyncio.sleep(5)  # 等待5秒后重试

    async def send_message(self, message):
        if self.websocket is not None:
            await self.websocket.send(message)

    async def receive_message(self):
        try:
            while True:
                message = await self.websocket.recv()
                if message == 'pong':
                    logging.info("ok")
                else:
                    handle_ws_message(message)
        except websockets.exceptions.ConnectionClosed:
            logging.info("Connection closed, attempting to reconnect")

    # 定期发送心跳或保持连接的请求
    async def keep_alive(self):
        while True:
            try:
                if self.websocket is not None:
                    await self.websocket.send("ping")
                    logging.info("Sent keep alive ping")
            except Exception as e:
                logging.info(f"Failed to send keep alive message: {e}")
            await asyncio.sleep(20)  # 每20秒发送一次


# 创建一个全局的 WebSocketClient 实例
websocket_client = WebSocketClient("wss://ws.bitget.com/v2/ws/public")


async def start_client():
    await websocket_client.connect()


def handle_ws_message(message):
    json_object = json.loads(message)
    if 'action' in json_object:
        if json_object['action'] == 'snapshot' or json_object['action'] == 'update':
            if json_object['arg']["channel"] == 'ticker':
                update_ticker(json_object)
            elif json_object['arg']["channel"] == 'candle1m':
                update_candle(json_object)


def update_ticker(json_object):
    """
    更新ticker数据
    :param message:
    :return:
    """
    if json_object['data'] is not None:
        bitgetTicker = BitgetTicker()
        bitgetTicker.ts = json_object['data'][0]['ts']
        bitgetTicker.create_time = datetime.datetime.fromtimestamp(float(bitgetTicker.ts) / 1000.0)
        bitgetTicker.last_price = json_object['data'][0]['lastPr']
        bitgetTicker.inst_id = json_object['data'][0]['instId']
        bitgetTicker.save()


def update_candle(json_object):
    '''
    更新k线
    :param json_object:
    :return:
    '''
    logging.info(json_object)
    if json_object['action'] == 'update':
        hedging_service.update_candle(json_object['arg']['instId'], json_object['data'][0])
        hedging_service.delete_candle()


def update_depth(json_object):
    """
    更新深度数据
    :param message:
    :return:
    """
    if json_object['data'] is not None:
        ts = json_object['data'][0]['ts']
        create_time = datetime.datetime.fromtimestamp(float(ts) / 1000.0)
        asks = json_object['data'][0]['asks']
        bids = json_object['data'][0]['bids']
        bitgetDepthList = []
        if len(asks) > 0:
            for index, item in enumerate(asks):
                bitgetDepth = BitgetDepth()
                bitgetDepth.inst_id = json_object['arg']['instId']
                if bitgetDepth.inst_id == "BTCUSDT":
                    bitgetDepth.currency_id = 1
                elif bitgetDepth.inst_id == "ETHUSDT":
                    bitgetDepth.currency_id = 2
                bitgetDepth.create_time = create_time
                bitgetDepth.direction = 'ask'
                bitgetDepth.ranking = index + 1
                bitgetDepth.price = item[0]
                bitgetDepth.sz = item[1]
                bitgetDepth.ts = ts
                bitgetDepthList.append(bitgetDepth)
        if len(bids) > 0:
            for index, item in enumerate(bids):
                bitgetDepth = BitgetDepth()
                bitgetDepth.inst_id = json_object['arg']['instId']
                if bitgetDepth.inst_id == "BTCUSDT":
                    bitgetDepth.currency_id = 1
                elif bitgetDepth.inst_id == "ETHUSDT":
                    bitgetDepth.currency_id = 2
                bitgetDepth.create_time = create_time
                bitgetDepth.direction = 'bid'
                bitgetDepth.ranking = index + 1
                bitgetDepth.price = item[0]
                bitgetDepth.sz = item[1]
                bitgetDepth.ts = ts
                bitgetDepthList.append(bitgetDepth)
        BitgetDepth.bulk_create(bitgetDepthList)
